import onnx
import json
import argparse
import sys

def main():
    parser = argparse.ArgumentParser(description="将 ONNX 模型转换为 JSON 格式")
    parser.add_argument("-i", "--input", required=True, help="输入的 ONNX 文件路径")
    parser.add_argument("-o", "--output", required=True, help="输出的 JSON 文件路径")
    
    args = parser.parse_args()
    
    try:
        # 加载 ONNX 模型
        model = onnx.load(args.input)
        graph = model.graph
        
        # 创建模型摘要
        summary = {
            "inputs": [{"name": i.name, "type": str(i.type)} for i in graph.input],
            "outputs": [{"name": o.name, "type": str(o.type)} for o in graph.output],
            "nodes": [
                {
                    "name": n.name,
                    "op_type": n.op_type,
                    "inputs": list(n.input),
                    "outputs": list(n.output),
                    "attributes": {a.name: str(a) for a in n.attribute}
                }
                for n in graph.node
            ]
        }
        
        # 写入 JSON 文件
        with open(args.output, "w", encoding="utf-8") as f:
            json.dump(summary, f, indent=2)
        
        print(f"成功将 {args.input} 转换为 {args.output}")
        
    except FileNotFoundError:
        print(f"错误：找不到文件 {args.input}")
        sys.exit(1)
    except Exception as e:
        print(f"错误：{str(e)}")
        sys.exit(1)

if __name__ == "__main__":
    main()